package com.example.batch.core.loadbalance;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.List;
import java.util.Map;
import java.util.ArrayList;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import java.sql.Connection;

@Component
@Slf4j
public class DataSourceLoadBalancer {
    
    private final Map<String, List<DataSourceWrapper>> dataSourceGroups = new ConcurrentHashMap<>();
    private final Map<String, LoadBalanceStrategy> strategies = new ConcurrentHashMap<>();
    
    @Autowired
    private DynamicRoutingDataSource dynamicDataSource;
    
    @PostConstruct
    public void init() {
        // 初始化数据源分组
        dynamicDataSource.getResolvedDataSources().forEach((name, dataSource) -> {
            String groupName = getGroupName(name.toString());
            dataSourceGroups.computeIfAbsent(groupName, k -> new ArrayList<>())
                .add(new DataSourceWrapper(name.toString(), dataSource));
        });
        
        // 初始化负载均衡策略
        dataSourceGroups.forEach((group, dataSources) -> {
            strategies.put(group, createStrategy(group));
        });
    }
    
    public DataSource getDataSource(String groupName) {
        List<DataSourceWrapper> availableDataSources = dataSourceGroups.get(groupName).stream()
            .filter(DataSourceWrapper::isAvailable)
            .collect(Collectors.toList());
            
        if (availableDataSources.isEmpty()) {
            throw new NoAvailableDataSourceException("No available data source in group: " + groupName);
        }
        
        LoadBalanceStrategy strategy = strategies.get(groupName);
        return strategy.select(availableDataSources);
    }
    
    @Scheduled(fixedRate = 5000)
    public void checkDataSourceStatus() {
        dataSourceGroups.values().stream()
            .flatMap(List::stream)
            .forEach(wrapper -> {
                try (Connection conn = wrapper.getDataSource().getConnection()) {
                    wrapper.setAvailable(conn.isValid(3000));
                    wrapper.setLastCheckTime(System.currentTimeMillis());
                } catch (Exception e) {
                    wrapper.setAvailable(false);
                    log.warn("Data source check failed: {}", wrapper.getName(), e);
                }
            });
    }
} 